// Copyright (c) Lawrence Livermore National Security, LLC and other VisIt
// Project developers.  See the top-level LICENSE file for dates and other
// details.  No copyright assignment is required to contribute to VisIt.

// ************************************************************************* //
//                       avtWavefrontOBJWriter.C                             //
// ************************************************************************* //

#include <avtWavefrontOBJWriter.h>

#include <vector>

#include <vtkAlgorithm.h>
#include <vtkAppendPolyData.h>
#include <vtkDataSet.h>
#include <vtkGeometryFilter.h>
#include <vtkDataSetWriter.h>
#include <vtkPolyDataReader.h>
#include <vtkPolyDataWriter.h>
#include <vtkCharArray.h>
#include <vtkPolyData.h>
#include <vtkTriangleFilter.h>
#include <DBOptionsAttributes.h>
#include <avtDatabaseMetaData.h>
#include <vtkImageData.h>
#include <vtkFloatArray.h>
#include <vtkPointData.h>
#include <avtColorTables.h>
#include <ColorTableAttributes.h>
#include <ColorControlPointList.h>
#include <ColorControlPoint.h>
#include <vtkOBJWriter.h>

#include <avtDatabaseMetaData.h>
#include <DebugStream.h>
#include <avtParallelContext.h>
#include <avtDatasetFileWriter.h>
#include <vtkPNGWriter.h>

using     std::string;
using     std::vector;


// ****************************************************************************
//  Method: avtWavefrontOBJWriter constructor
//
//  Programmer: pugmire -- generated by xml2avt
//  Creation:   Tue Apr 16 08:33:21 PDT 2013
// 
//  Modifications:
//    Justin Privitera, Fri Nov  3 15:25:32 PDT 2023
//    Now takes options.
// 
//    Justin Privitera, Tue Nov 28 17:31:40 PST 2023
//    Grab whether or not we want the color table inverted.
//
// ****************************************************************************

avtWavefrontOBJWriter::avtWavefrontOBJWriter(const DBOptionsAttributes *atts)
{
    doColor = atts->GetBool("Output colors");
    colorTable = atts->GetString("Color table");
    invertCT = atts->GetBool("Invert color table");
}

// ****************************************************************************
//  Method: avtWavefrontOBJWriter::OpenFile
//
//  Purpose:
//      Does no actual work.  Just records the stem name for the files.
//
//  Programmer: pugmire -- generated by xml2avt
//  Creation:   Tue Apr 16 08:33:21 PDT 2013
//
// ****************************************************************************

void
avtWavefrontOBJWriter::OpenFile(const string &stemname, int numblocks)
{
    stem = stemname;
}


// ****************************************************************************
//  Method: avtWavefrontOBJWriter::WriteHeaders
//
//  Purpose:
//      Writes out a VisIt file to tie the WavefrontOBJ files together.
//
//  Programmer: pugmire -- generated by xml2avt
//  Creation:   Tue Apr 16 08:33:21 PDT 2013
//
// ****************************************************************************

void
avtWavefrontOBJWriter::WriteHeaders(const avtDatabaseMetaData *md,
                                    const vector<string> &scalars,
                                    const vector<string> &vectors,
                                    const vector<string> &materials)
{
    // WRITE OUT HEADER INFO
}


// ****************************************************************************
//  Method: avtWavefrontOBJWriter::WriteChunk
//
//  Purpose:
//      This writes out one chunk of an avtDataset.
//
//  Programmer: pugmire -- generated by xml2avt
//  Creation:   Tue Apr 16 08:33:21 PDT 2013
//
//  Modifications:
//    Brad Whitlock, Tue Sep  8 17:04:07 PDT 2015
//    Rely on base class for geometry consolidation. This method will only
//    be called on group leaders.
// 
//    Justin Privitera, Fri Nov  3 15:25:32 PDT 2023
//    Added ability to write out mtllib and texture.
//
// ****************************************************************************

void
avtWavefrontOBJWriter::WriteChunk(vtkDataSet *ds, int chunk)
{
    std::string filename;
    if(writeContext.GroupSize() > 1)
    {
        char ext[20];
        snprintf(ext, 20, ".%d.obj", writeContext.GroupRank());
        filename = stem + ext;
    }
    else
        filename = stem + ".obj";

    if (doColor)
    {
        std::string textureFilename = stem + ".png";

        avtDatasetFileWriter::WriteOBJFile(ds, // dataset to write
                                           filename.c_str(), // filename
                                           nullptr, // label
                                           true, // YES writeMTL
                                           true, // YES MTLHasTex
                                           textureFilename); // name of texture file

        vtkImageData *image = GetColorTable();
        vtkImageWriter *writer = vtkPNGWriter::New();
        writer->SetFileName(textureFilename.c_str());
        writer->SetInputData(image);
        writer->Write();
        writer->Delete();
    }
    else
    {
        avtDatasetFileWriter::WriteOBJFile(ds, // dataset to write
                                           filename.c_str(), // filename
                                           nullptr); // label
    }
}

// ****************************************************************************
//  Method: avtWavefrontOBJWriter::CloseFile
//
//  Purpose:
//      Closes the file.  This does nothing in this case.
//
//  Programmer: pugmire -- generated by xml2avt
//  Creation:   Tue Apr 16 08:33:21 PDT 2013
//
// ****************************************************************************

void
avtWavefrontOBJWriter::CloseFile(void)
{
}

// ****************************************************************************
// Method: avtWavefrontOBJWriter::CreateTrianglePolyData
//
// Purpose:
//   Tell VisIt's export that we'll want triangles.
//
// Returns:    True
//
// Programmer: Brad Whitlock
// Creation:   Tue Sep  8 17:00:23 PDT 2015
//
// Modifications:
//
// ****************************************************************************

bool
avtWavefrontOBJWriter::CreateTrianglePolyData() const
{
    return true;
}

// ****************************************************************************
//  Method: avtWavefrontOBJWriter::GetCombineMode
//
//  Purpose:
//     Provides a hint to the export mechanism to tell it how to combine data.
//
//  Note: We combine geometry because WavefrontOBJ tools will want 1 file.
//
//  Programmer: Brad Whitlock
//  Creation:   Tue Sep  8 15:36:45 PDT 2015
//
// ****************************************************************************

avtDatabaseWriter::CombineMode
avtWavefrontOBJWriter::GetCombineMode(const std::string &) const
{
    return CombineAll;
}

//****************************************************************************
// Method:  avtWavefrontOBJWriter::GetColorTable()
//
// Purpose:
//   This function is creating a vtkImageData which contains a 1D strip of
//   pixels representing the selected color table. A vtkImageData is needed
//   to present to the vtkPNGWriter so that it can write out the texture
//   for the OBJ file.
//
// Programmer:  Justin Privitera
// Creation:    11/03/23
//
// Modifications:
//    Justin Privitera, Mon Nov 27 14:57:17 PST 2023
//    I added extra pixels on either end of the color table to act as padding
//    so that texture coordinates do not fall off the ends (and wrap around).
// 
//    Justin Privitera, Tue Nov 28 17:31:40 PST 2023
//    We now handle the inverted color table case.
//
//****************************************************************************

vtkImageData *
avtWavefrontOBJWriter::GetColorTable()
{
    const ColorTableAttributes *colorTables = avtColorTables::Instance()->GetColorTables();
    const ColorControlPointList *table = colorTables->GetColorControlPoints(colorTable);

    if (! table)
    {
        return nullptr;
    }
    

    // We don't have color tables that have this many control points,
    // so this should be a good choice for the number of colors.
    const int ncolors = 256;
    unsigned char rgb[ncolors * 3];
    
    table->GetColors(rgb, ncolors);

    vtkImageData *imageData = vtkImageData::New();

    // We need two extra colors, so (ncolors - 1 + 2) -> (ncolors + 1)
    // hence x: [0, ncolors + 1], y: [0, 0], z: [0, 0]
    // These pixels will sit on the ends and act as padding so the texture coords
    // don't lead to strange behavior with max and min values wrapping around
    imageData->SetExtent(0, ncolors + 1, 0, 0, 0, 0);
    imageData->SetSpacing(1., 1., 1.);
    imageData->SetOrigin(0., 0., 0.);
    imageData->AllocateScalars(VTK_UNSIGNED_CHAR, 3);
    unsigned char *pixels = (unsigned char *)imageData->GetScalarPointer(0, 0, 0);
    unsigned char *ipixel = pixels;

    if (invertCT)
    {
        // the first extra pixel will get the same color as the first real pixel
        int last_index = ncolors * 3 - 3;
        *ipixel = rgb[last_index];
        ipixel ++;
        *ipixel = rgb[last_index + 1];
        ipixel ++;
        *ipixel = rgb[last_index + 2];
        ipixel ++;

        // iterate through the colors in reverse
        for (int i = ncolors * 3 - 3; i >= 0; i -= 3)
        {
            *ipixel = rgb[i];
            ipixel ++;
            *ipixel = rgb[i + 1];
            ipixel ++;
            *ipixel = rgb[i + 2];
            ipixel ++;
        }

        // the second (and last) extra pixel will get the same color as the last real pixel
        *ipixel = rgb[0];
        ipixel ++;
        *ipixel = rgb[1];
        ipixel ++;
        *ipixel = rgb[2];
        ipixel ++;
    }
    else
    {
        // the first extra pixel will get the same color as the first real pixel
        *ipixel = rgb[0];
        ipixel ++;
        *ipixel = rgb[1];
        ipixel ++;
        *ipixel = rgb[2];
        ipixel ++;

        // iterate through the colors
        for (int i = 0; i < ncolors * 3; i += 3)
        {
            *ipixel = rgb[i];
            ipixel ++;
            *ipixel = rgb[i + 1];
            ipixel ++;
            *ipixel = rgb[i + 2];
            ipixel ++;
        }

        // the second (and last) extra pixel will get the same color as the last real pixel
        int last_index = ncolors * 3 - 3;
        *ipixel = rgb[last_index];
        ipixel ++;
        *ipixel = rgb[last_index + 1];
        ipixel ++;
        *ipixel = rgb[last_index + 2];
        ipixel ++;
    }

    return imageData;
}
